from IPython import display
from jqdatasdk.technical_analysis import MACD, TRIX
from jqdatasdk import finance, query, indicator, auth, valuation, income
from jqdatasdk.api import attribute_history_engine, get_fundamentals, get_industry_stocks, get_security_info, get_all_securities
import pandas as pd
import numpy as np
import datetime
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['font.sans-serif'] = ['KaiTi']
mpl.rcParams['font.serif'] = ['KaiTi']
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题,或者转换负号为字符串

auth('你的聚宽账号', '你的密码')

this_year = datetime.datetime.now().year-1
this_month = datetime.datetime.now().month
this_quarter = (this_month - 1) // 3 + 1


def roe_select(codes, count=10):
    """
    最近10年平均ROE大于20%，并且每年ROE大于15%。
    净资产收益率ROE(%)=归属于母公司股东的净利润*2/（期初归属于母公司股东的净资产+期末归属于母公司股东的净资产）
    """
    # 找出已经上市十年了的代码
    ten_year_codes = []
    df = finance.run_query(query(finance.STK_LIST.code, finance.STK_LIST.start_date).filter(
        finance.STK_LIST.code.in_(codes)))
    for row in df.itertuples():
        if getattr(row, 'start_date').year <= this_year - count:
            ten_year_codes.append(getattr(row, 'code'))

    roe_selected = []
    for code in ten_year_codes:
        q = query(
            indicator.statDate, indicator.roe
        ).filter(
            indicator.code == code
        )
        rets = [get_fundamentals(q, statDate=str(i))
                for i in range(this_year - count, this_year)]
        df = pd.concat(rets)
        df.set_index('statDate', inplace=True)
        sum = 0
        for row in df.itertuples():
            roe = getattr(row, 'roe')
            if roe < 15:
                break
            sum += roe
        else:
            if sum/count > 20:
                roe_selected.append(code)
    return roe_selected


def industry_average_pe_ratio(industry_code, date=None, stateDate=None):
    """
    行业平均市盈率
    industry_code: 行业代码
    """
    codes = get_industry_stocks(industry_code)
    df = get_fundamentals(query(
        # pe_ratio市盈率(PE, TTM)每股市价为每股收益的倍数，反映投资人对每元净利润所愿支付的价格，用来估计股票的投资报酬和风险
        valuation.code, valuation.pe_ratio
    ).filter(
        valuation.code.in_(codes)
    ), date=date, statDate=stateDate)
    sum = 0
    for row in df.itertuples():
        pe_ratio = getattr(row, 'pe_ratio')
        sum += pe_ratio
    return sum / len(codes)


def show_pe_ratio_roe(codes, count=10):
    """
    显示指定股票的pe和roe
    count: 从现在往回找几年
    """
    data = {}
    codes_name = ''
    pe_ratio_df = pd.DataFrame(
        index=[i for i in range(this_year - count, this_year)])
    roe_df = pd.DataFrame(
        index=[i for i in range(this_year - count, this_year)])
    for code in codes:
        data[code] = []
        codes_name += get_security_info(code).display_name + '、'
        que = query(
            valuation.pe_ratio, indicator.roe
        ).filter(
            valuation.code == code
        )
        the_df = pd.concat([get_fundamentals(que, statDate=str(i))
                            for i in range(this_year - count, this_year)])
        pe_ratio_df.insert(loc=0, column=code, value=the_df.pe_ratio.to_list())
        roe_df.insert(loc=0, column=code, value=the_df.roe.to_list())
    codes_name = codes_name[0:-1]
    pe_ratio_df.plot(title=codes_name + '历史PE')
    plt.show()
    roe_df.plot(title=codes_name + '历史roe')
    plt.show()


# 技术指标https://www.joinquant.com/help/api/help?name=technicalanalysis

# 利润apihttps://www.joinquant.com/help/api/help?name=Stock#%E5%88%A9%E6%B6%A6%E6%95%B0%E6%8D%AE

def get_cpg(code, count=5):
    """
    复合利润增长率=(当前净利润/基础净利润)^(1/时间) - 1
    """
    q = query(
        income.net_profit
    ).filter(
        income.code == code
    )
    df_1 = get_fundamentals(q, statDate=str(this_year-count))
    df_2 = get_fundamentals(q, statDate=str(this_year-1))
    try:

        return (df_2.iat[0, 0]/df_1.iat[0, 0])**(1/count)-1
    except IndexError:
        print(code)


def show_cpg_select(codes, count=5, limit=10):
    """
    最近5年内复合利润增长率＞10
    """
    data = {}
    df = None
    for code in codes:
        profit_rate = get_cpg(code, count)
        try:
            if profit_rate*100 > limit:
                data[code] = profit_rate
        except TypeError:
            continue
    df = pd.DataFrame(data.values(), index=data.keys(),
                      columns=['复合利润增长率'])
    df.plot.bar(title='复合利润增长率')
    plt.show()
    return [*data.keys()]


get_cpg('')
